import numpy as np
import logging
from scipy.misc import imread
from scipy.misc import imsave
from scipy.signal import convolve2d
from scipy.stats import hmean
from math import floor


# Parameters
# ----------
# image : ndarray
#     Input image data. Will be converted to float. Must be gray scale image
# mode : str
#     One of the following strings, selecting the type of noise to add:
#
#     'gauss'     Gaussian-distributed additive noise.
#     'poisson'   Poisson-distributed noise generated from the data.
#     's&p'       Replaces random pixels with 0 or 1.
#     'speckle'   Multiplicative noise using out = image + n*image,where
#                 n is uniform noise with specified mean & variance.
#
def noisy(image, noise_typ, mean=0, var=0.02, amount=0.1, s_vs_p=0.5):
    if noise_typ == "gauss":
        row, col = image.shape
        # mean = 0
        # var = 0.1
        sigma = var ** 0.5
        gauss = np.random.normal(mean, sigma, (row, col))
        gauss = gauss.reshape(row, col)
        noisy = image + gauss
        noisy[noisy > 1] = 1
        noisy[noisy < 0] = 0
        return noisy
    elif noise_typ == "s&p":
        # s_vs_p = 0.5
        # amount = 0.004
        out = np.copy(image)
        # Salt mode
        num_salt = np.ceil(amount * image.size * s_vs_p)
        coords = [np.random.randint(0, i - 1, int(num_salt)) for i in image.shape]
        out[coords] = 1

        # Pepper mode
        num_pepper = np.ceil(amount * image.size * (1. - s_vs_p))
        coords = [np.random.randint(0, i - 1, int(num_pepper)) for i in image.shape]
        out[coords] = 0
        return out
    elif noise_typ == "poisson":
        vals = len(np.unique(image))
        vals = 2 ** np.ceil(np.log2(vals))
        noisy = np.random.poisson(image * vals) / float(vals)
        noisy[noisy > 1] = 1
        noisy[noisy < 0] = 0
        return noisy
    elif noise_typ == "speckle":
        row, col = image.shape
        gauss = np.random.randn(row, col)
        gauss = gauss.reshape(row, col)
        noisy = image + image * gauss
        noisy[noisy > 1] = 1
        noisy[noisy < 0] = 0
        return noisy
    else:
        logging.log(logging.ERROR, "Noise Type Not Found!")


def meanFilter(img, type, shape):
    assert  shape[0] == shape[1]
    size = floor(shape[0] / 2)
    if type == 'ari':
        return convolve2d(img, np.ones(shape) / shape[0] / shape[0], 'same', 'symm')
    elif type == 'geo':
        oimg = img.copy()
        for row in range(size,img.shape[0]-size):
            for col in range(size, img.shape[1]-size):
                oimg[row,col] = \
                    np.prod(img[row-size:row+size+1, col-size:col+size+1]) \
                    ** (1/shape[0]/shape[1])
        return oimg
    elif type == 'har':
        img = img.copy() + 1/255
        oimg = img.copy()
        for row in range(size, img.shape[0] - size):
            for col in range(size, img.shape[1] - size):
                patch = img[row - size:row + size + 1, col - size:col + size + 1]
                patch = np.reshape(patch, (shape[0] * shape[1], 1))
                v1 = hmean(patch)
                oimg[row, col] = v1
        return oimg
    elif type == 'con':
        q = 4
        img = img.copy() + 1/255
        oimg = img.copy()
        for row in range(size, img.shape[0] - size):
            for col in range(size, img.shape[1] - size):
                patch = img[row - size:row + size + 1, col - size:col + size + 1]
                oimg[row, col] = np.sum(np.power(patch,q+1)) / np.sum(np.power(patch,q))
        return oimg
    else:
        assert False


def main():
    FORMAT = '	[%(levelname)-5s]%(asctime)-8s %(filename)s:%(lineno)d %(message)s'
    DATEFORMAT = '%H:%M:%S'
    logging.basicConfig(level=logging.DEBUG, format=FORMAT, datefmt=DATEFORMAT)

    def applyFilter(img, name, imgName):
        oimg = meanFilter(img,name,(3,3))
        imsave('out/hw4_%s_noise.jpg' % (imgName), img)
        imsave('out/hw4_%s_noise_%s_filter.jpg' % (imgName, name), oimg)

    image = imread('data/Circuit.tif') / 255
    noiseGaussion = noisy(image, 'gauss')
    noiseS = noisy(image, 's&p', s_vs_p=1)
    noiseP = noisy(image, 's&p', s_vs_p=0)

    applyFilter(noiseGaussion,'ari', 'gaussian')
    applyFilter(noiseGaussion,'geo', 'gaussian')
    applyFilter(noiseGaussion,'har', 'gaussian')
    applyFilter(noiseGaussion,'con', 'gaussian')

    applyFilter(noiseP,'ari', 'pepper')
    applyFilter(noiseP,'geo', 'pepper')
    applyFilter(noiseP,'har', 'pepper')
    applyFilter(noiseP,'con', 'pepper')

    applyFilter(noiseS,'ari', 'salt')
    applyFilter(noiseS,'geo', 'salt')
    applyFilter(noiseS,'har', 'salt')
    applyFilter(noiseS,'con', 'salt')



if __name__ == '__main__':
    main()
